import cv2
import numpy as np


def detect_ellipse(hsvimg, color_low, color_high):
    thres = cv2.inRange(hsvimg, color_low, color_high)
    median_filter = cv2.medianBlur(thres, 7)  # 中值滤波
    kernal = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    mask = cv2.morphologyEx(median_filter, cv2.MORPH_CLOSE, kernal)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernal)
    contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    contours = contours[0]
    
    ellipse_in_left = None
    ellipse_in_right = None
    for cnt in contours:
        if cnt.size < 200: # 轮廓中点的数量
            print('cnt size')
            print(cnt.size)
            continue
        cnt_area = cv2.contourArea(cnt)
        if cnt_area < 1000: # 轮廓围成的点的面积
            print('cnt area')
            print(cnt_area)
            continue
        # # 拟合圆形
        # (x,y),radius = cv2.minEnclosingCircle(cnt)
        # center = (int(x),int(y))
        # radius = int(radius)
        # img = cv2.circle(img,center,radius,(0,255,0),2)
        # # 拟合长方形
        # x, y, w, h = cv2.boundingRect(cnt)
        # cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 255), 2)
        # 拟合椭圆
        ellipse = cv2.fitEllipse(cnt)
        ellipse_size =  ellipse[1]
        if ellipse_size[1]/ellipse_size[0] < 0.8 or ellipse_size[1]/ellipse_size[0] > 1.2:
            print(ellipse_size[1]/ellipse_size[0])
            continue
        ellipse_center = ellipse[0]

        if ellipse_center[0] < hsvimg.shape[1]/2:
            ellipse_in_left = ellipse
        else:
            ellipse_in_right = ellipse

        if ellipse_in_left is not None and ellipse_in_right is not None:
            return ellipse_in_left, ellipse_in_right # 返回第一个椭圆
    return None, None


def triangulation(left_point, right_point, K_l, D_l, K_r, D_r, P, R):
    left_undistorted_point = cv2.undistortPoints(left_point, K_l, D_l)
    right_undistorted_point = cv2.undistortPoints(right_point, K_r, D_r)
    left_undistorted_point_crossmat = np.zeros((3,3))
    left_undistorted_point_crossmat[0,1] = -1
    left_undistorted_point_crossmat[0,2] = left_undistorted_point[0,0,1]
    left_undistorted_point_crossmat[1,0] = 1
    left_undistorted_point_crossmat[1,2] = -left_undistorted_point[0,0,0]
    left_undistorted_point_crossmat[2,0] = -left_undistorted_point[0,0,1]
    left_undistorted_point_crossmat[2,1] = left_undistorted_point[0,0,0]

    right_undistorted_point_crossmat = np.zeros((3,3))
    right_undistorted_point_crossmat[0,1] = -1
    right_undistorted_point_crossmat[0,2] = right_undistorted_point[0,0,1]
    right_undistorted_point_crossmat[1,0] = 1
    right_undistorted_point_crossmat[1,2] = -right_undistorted_point[0,0,0]
    right_undistorted_point_crossmat[2,0] = -right_undistorted_point[0,0,1]
    right_undistorted_point_crossmat[2,1] = right_undistorted_point[0,0,0]

    T0 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0]])
    T_L_R = np.concatenate([R.transpose(), P], 1) # 从L到R的转换矩阵

    A = np.concatenate([np.matmul(left_undistorted_point_crossmat,T0), np.matmul(right_undistorted_point_crossmat, T_L_R)], 0)
    U, S, V = np.linalg.svd(A)

    if V[3,3] < 1e-6 and  V[3,3] > -1e-6:
        return None

    point = V[3,:].transpose()/V[3,3]

    return point

def undistort_image(distorted_image, K, D):
    undistorted_image = cv2.undistort(distorted_image, K, D)
    cv2.imshow('distort', distorted_image)
    cv2.imshow('undistort', undistorted_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()